import pandas as pd
import numpy as np
import os
from LMdata.LMdataset import idx2kp
from utils.metrix import cal_NEscore

files = [
    'GCN336_bs4-5.85.csv',
    'GCN512_bs4-5.25.csv',
]

temp = pd.read_csv(os.path.join('./pred_online1/csvs', files[0]))
ori_ids = temp['image_id'].tolist()
ori_cates = temp['img_category'].tolist()
vis_mat = np.zeros((temp.shape[0], 24))
for i in range(24):
    vis_mat[:,i] = temp[idx2kp[i]].str.split('_').str[2].astype(int)  # vis  [-1, 0, 1]



pred_xyv = np.zeros((len(files), temp.shape[0], 24, 3), dtype=np.int)
for j,file in enumerate(files):
    df = pd.read_csv(os.path.join('./pred_online1/csvs', file))

    for i in range(24):
        pred_xyv[j, :, i, 0] = df[idx2kp[i]].str.split('_').str[0].astype(int)  # x coord [-1,512]
        pred_xyv[j, :, i, 1] = df[idx2kp[i]].str.split('_').str[1].astype(int)  # y coord [-1,512]


pred_xyv = pred_xyv.mean(0).astype(int)
pred_xyv[:,:,2] = vis_mat

# generate sub file
sub = pd.DataFrame()
sub['image_id'] = ori_ids
sub['img_category'] = ori_cates

for i in range(24):
    kp_str = idx2kp[i]
    sub[kp_str] = pred_xyv[:, i, 0].astype(str)
    sub[kp_str] = sub[kp_str].str[:] + \
                  '_' + pred_xyv[:, i, 1].astype(str) + \
                  '_' + pred_xyv[:, i, 2].astype(str)

    sub[kp_str][pred_xyv[:, i, 2] == -1] = '-1_-1_-1'

sub.to_csv('./pred_online1/csvs/avg1.csv',index=False)

